/*
Copyright 2013 Edwin Hoeks All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

Author: ehoeks@gmail.com (Edwin Hoeks)
*/

#include "png_optimizer.h"

#include <stdio.h>

#include "zlib_container.h" /* zopfli zlib deflate */
#include "png.h" /* libpng */


png_byte png_sig[8] = {137, 80, 78, 71, 13, 10, 26, 10};

png_byte png_IDAT[5] = { 73,  68,  65,  84, '\0'};
png_byte png_IEND[5] = { 73,  69,  78,  68, '\0'};


/*
crc = crc32(0, Z_NULL, 0); // reset CRC
crc = crc32(crc, ptr, length); // update crc
png_save_uint_32(buffer, crc); // save crc to buffer[4]
*/

typedef struct ImageData {

    /*  */
    unsigned char* source;
    unsigned int source_len;
    unsigned char* source_in;
    unsigned int source_left;
    
    /*  */
    unsigned char* raw;
    unsigned long int raw_len;
    
    /*  */
    unsigned char* dest;
    size_t dest_len;
    
} ImageData;
/*
default to 4 Mbyte source, 16 Mbyte raw
*/
void InitImageData2(ImageData* imagedata, unsigned int size) {
    /* quarter sized source buffer */
    imagedata->source_len = size / 4;
    imagedata->source_left = imagedata->source_len;
    imagedata->source = malloc(imagedata->source_len);
    imagedata->source_in = imagedata->source;
    
    /* full size raw buffer */
    imagedata->raw_len = size;
    imagedata->raw = malloc(imagedata->raw_len);
    
    /* empty destination buffer */
    imagedata->dest_len = 0;
    imagedata->dest = 0;
    
    
}

void InitImageData(ImageData* imagedata) {
    InitImageData2(imagedata, (16 * 1024 * 1024));
}

void AddSourceImageData(ImageData* imagedata,
    const unsigned char* data,
    size_t length) { 
    if (length > imagedata->source_left) {
        fprintf(stderr, "compressed image data does not fit buffer\n");
        exit(-1);
    } else {
        imagedata->source_left -= length;
        memcpy(imagedata->source_in, data, length);
        imagedata->source_in += length;
    }
}

int IsImageDataFilled(ImageData* imagedata) {
    return (imagedata->source_len != imagedata->source_left);
}

/*
ZEXTERN int ZEXPORT uncompress OF((Bytef *dest,   uLongf *destLen,
                                   const Bytef *source, uLong sourceLen));
*/
void RecompressImageData(ImageData* imagedata,const Options* options)
{
    int res;
  
    /* UnCompress */
    res = uncompress(imagedata->raw, &(imagedata->raw_len),
        imagedata->source, (imagedata->source_len - imagedata->source_left));
    
    if (Z_OK == res) {
        fprintf(stderr, "uncompress succes! (%u -> %u)\n", 
            (imagedata->source_len - imagedata->source_left),
            (unsigned int)imagedata->raw_len);
    } else if (Z_MEM_ERROR == res) {
        fprintf(stderr, "uncompress returned Z_MEM_ERROR\n");
        exit(-1);
    } else if (Z_BUF_ERROR == res) {
        fprintf(stderr, "uncompress returned Z_BUF_ERROR\n");
        exit(-1);
    } else if (Z_DATA_ERROR == res) {
        fprintf(stderr, "uncompress returned Z_DATA_ERROR\n");
        exit(-1);
    } else {
        fprintf(stderr, "uncompress returned an unspecified error: %d\n", res);
        exit(-1);
    }
    
    /* Recompress using zopfli */
    ZlibCompress(options,
        imagedata->raw, imagedata->raw_len,
        &(imagedata->dest), &(imagedata->dest_len));
        
    fprintf(stderr, "compress succes! (%u -> %u)\n", 
        (imagedata->source_len - imagedata->source_left),
        (unsigned int)imagedata->dest_len);
}

void FreeImageData(ImageData* imagedata) {
    /* source */
    if (imagedata->source) {
        free(imagedata->source);
        imagedata->source = 0;
        imagedata->source_len = 0;
    }
    /* raw */
    if (imagedata->raw) {
        free(imagedata->raw);
        imagedata->raw = 0;
        imagedata->raw_len = 0;
    }
    /* dest */
    if (imagedata->dest) {
        free(imagedata->dest);
        imagedata->dest = 0;
        imagedata->dest_len = 0;
    }
}


/*
Recompresses according to the png specification.
*/
void PngRecompress(const Options* options,
                  const unsigned char* in, size_t insize,
                  unsigned char** out, size_t* outsize) {

    png_bytep in_png = (png_bytep)in;
    /* Check png signature */
    if ((insize < 8) ||
        (0 != png_sig_cmp(in_png, 0, 8))) {
        fprintf(stderr, "Input file is not a valid png file (signature check failed)\n");
        exit(-1);
    } else {
        /* write png signature */
        unsigned int sig_index;
        for(sig_index = 0; sig_index < sizeof(png_sig); ++sig_index) {
            APPEND_DATA(png_sig[sig_index], out, outsize);
        }
    }
    insize -= 8;
    in_png +=8;
    
    if (insize < 8) {
        fprintf(stderr, "Input file is not a valid png file (content missing)\n");
        exit(-1);
        
    } else {
        /* flag to indicate the infile is read completely
            (IEND and insize ==0) */
        int finished = 0;
        ImageData imagedata;
        
        InitImageData(&imagedata);
        
        /* process chunk by chunk */
        do {
            /* flag to indicate if this chunck must be passsed through or
                processed */
            int passthrough = 1;
            
            /* get the length of the chunk */
            png_uint_32 chunk_data_length = png_get_uint_32(in_png);
            /* get the chunk type */
            png_byte chunk_name[4];
            memcpy(chunk_name, in_png + 4, 4);
            
            if (!memcmp(chunk_name, png_IDAT, 4)) {
                passthrough = 0;
                
                fprintf(stderr, "IDAT %u bytes %u Kbytes\n", (unsigned int)chunk_data_length, ((unsigned int)chunk_data_length/1024));
                
                AddSourceImageData(&imagedata, in_png + 8, chunk_data_length);
                
                insize -= 8 + chunk_data_length + 4;
                in_png += 8 + chunk_data_length + 4;
            }
            
            /* check for IEND */
            if (!memcmp(chunk_name, png_IEND, 4)) {
                finished = 1;
            }
            
            /* recompress iamge data just before the next chuck after the image
            data chunks is passed through */
            if (passthrough && IsImageDataFilled(&imagedata)) {
                png_uint_32 crc;
                png_byte buffer[4];\
                unsigned int index;
                
                /* RecompressImageData */
                RecompressImageData(&imagedata, options);
                
                /* calculate CRC */
                crc = crc32(0, Z_NULL, 0);
                crc = crc32(crc, png_IDAT, 4); /* IDAT */
                crc = crc32(crc, imagedata.dest, imagedata.dest_len); /* compressed data (payload)*/
                
                /* Send the lenght */
                png_save_uint_32(buffer, (unsigned int)imagedata.dest_len);
                for (index = 0; index < sizeof(buffer); index++) {
                    APPEND_DATA(buffer[index], out, outsize);
                }
                
                /* Send the chuck type */
                for (index = 0; index < sizeof(buffer); index++) {
                    APPEND_DATA(png_IDAT[index], out, outsize);
                }
                
                /* Send reprocessed image data */
                for (index = 0; index < imagedata.dest_len; index++) {
                    APPEND_DATA(imagedata.dest[index], out, outsize);
                }
                
                /* Send the crc */
                png_save_uint_32(buffer, crc);
                for (index = 0; index < sizeof(buffer); index++) {
                    APPEND_DATA(buffer[index], out, outsize);
                }
                
                FreeImageData(&imagedata);
                InitImageData(&imagedata);
            }
            
            
            if (passthrough) {
                const unsigned int bytes_to_passthrough = 8 + chunk_data_length + 4;
                unsigned int passthrough_index;
                for (passthrough_index = 0;
                    passthrough_index < bytes_to_passthrough;
                    passthrough_index++) {
                    /* copy byte for byte */
                    APPEND_DATA(in_png[passthrough_index], out, outsize);
                }
                
                insize -= bytes_to_passthrough;
                in_png += bytes_to_passthrough;
            }
            
        } while (!finished);
        
        FreeImageData(&imagedata);
    }
    
    
    
}

